# stdlib
require 'openssl'

# chaffing-rb
require 'random'
require 'binstr'

module Chaffing
  #
  # for a particular algorithm (i.e. OpenSSL::Digest::Digest) and MAC key, 
  # validate and generate MACs for any given payload
  #
  class MACr
    attr_accessor :key

    # returns a binary MAC value
    #
    def self.generate(digest, key, payload)
      OpenSSL::HMAC.digest(digest, key, payload)
    end

    def self.valid?(digest, key, payload, mac)
      mac == generate(digest, key, payload)
    end

    def initialize(algorithm, key)
      self.algorithm = algorithm
      self.key = key
    end

    def algorithm
      @digest.name
    end

    def algorithm=(input)
      @digest = OpenSSL::Digest::Digest.new(input)
    end

    def generate(payload)
      self.class.generate(@digest, @key, payload)
    end

    def valid?(payload, mac)
      self.class.valid?(@digest, @key, payload, mac)
    end
  end

  #
  # the abstract representation of blocks, as described in the '98 Rivest paper
  #
  class Block
    attr_accessor :serial, :payload, :mac

    def initialize(serial, payload, mac)
      @serial = serial
      @payload = payload
      @mac = mac
    end

    def to_a
      [@serial, @payload, @mac]
    end

    def to_h
      { :serial => @serial, :payload => @payload, :mac => @mac }
    end

    def to_s
      to_a.join
    end
  end

  #
  # for a particular MACr (i.e. algorithm and key), separate the wheat from 
  # chaff given any array of blocks.  also, process the wheat blocks to provide
  # the original text
  #
  class Winnower
    def initialize(macr)
      @macr = macr
    end

    def winnow(blocks)
      blocks.reject { |block| !@macr.valid?(block.payload, block.mac) }
    end

    def process(blocks)
      winnow(blocks).sort { |b1, b2| b1.serial <=> b2.serial }.
        map { |b| b.payload }.join
    end
  end


  #
  # for a particular MACr, generate an array of chaffed blocks for any given
  # text
  #
  class Chaffer

    # given an input text, create wheat blocks
    #
    def self.make_blocks(text, macr, size = 40)
      blocks = []
      tmp = text.dup
      count = 0
      while !tmp.empty?
        block = tmp.slice!(0, size)
        blocks << Block.new(count, block, macr.generate(block))
        count += 1
      end
      blocks
    end

    # given a single wheat block, add =chaff_per= chaff blocks and shuffle
    #
    def self.chaff_one_wheat(wheat_block, chaff_per)
      chaffed = [wheat_block]
      chaff_per.times { 
        chaffed << Block.new(wheat_block.serial,
                             Random::pseudo_bytes(wheat_block.payload.length),
                             Random::pseudo_bytes(wheat_block.mac.length))
      }
      chaffed.sort_by { rand }
    end
    
    def self.chaff_one_wheat_inject(wheat_block, chaff_per)
      Array.new(chaff_per).inject([wheat_block]) { |memo, dummy|
        memo + [Block.new(wheat_block.serial,
                          Random::pseudo_bytes(wheat_block.payload.length),
                          Random::pseudo_bytes(wheat_block.mac.length))]
      }.sort_by { rand }
    end
    
    # does not shuffle -- 1,1,1,2,2,2, ... n-1,n-1,n-1,n,n,n
    def self.chaff_several_wheat(wheat_blocks, chaff_per)
      wheat_blocks.inject([]) { |memo, wb| 
        memo + chaff_one_wheat(wb, chaff_per) 
      }
    end

    attr_accessor(:macr, :block_size, :chaff_per)

    def initialize(macr, block_size = 1, chaff_per = 1)
      @macr = macr
      @block_size = block_size
      @chaff_per = chaff_per
    end

    def mac_key
      @macr.key
    end

    def mac_key=(input)
      @macr.key = input
    end

    def mac_algorithm
      @macr.algorithm
    end

    def mac_algorithm=(input)
      @macr.algorithm = input
    end

    # returns an array of blocks
    #
    def chaffed(text)
      self.class.chaff_several_wheat(self.class.make_blocks(text, 
                                                            @macr, 
                                                            @block_size),
                                     @chaff_per)
    end
  end
end

if __FILE__ == $0
  def full_line_label(str)
    max_len = 80
    fillchar = '='
    label = ' ' << str.slice(0, max_len - 10) << ' '
    extra = label.length % 2
    margin = ((max_len - label.length) / 2) - 1
    [fillchar * (margin), label, fillchar * (margin + extra)].join(' ')
  end

  MAC_ALGO = 'sha'
  MAC_KEY = 'SICP'

  ## shared between chaffer and winnower
  macr = Chaffing::MACr.new(MAC_ALGO, MAC_KEY)
  ##


  text = 'Programs must be written for people to read, and only incidentally for machines to execute.'

  include Chaffing
  c = Chaffer.new(macr)
  w = Winnower.new(macr)

  # create the chaffed blocks
  chaffed_blocks = c.chaffed(text)

  # serialize the blocks and send them over an imaginary communications channel

  # here is the payload, naively reconstructed by a listener that doesn't know 
  # the MAC key
  chaffed_payload = chaffed_blocks.map { |b| b.payload }.join

  # here is the payload, smartly reconstructed by a listener that knows the 
  # MAC key
  winnowed_payload = w.process(chaffed_blocks)

  puts [full_line_label("text"), text].join("\n")
  puts
  puts [full_line_label("chaffed"), BinStr.display(chaffed_payload),
        full_line_label("chaffed hex"), BinStr.bytes(chaffed_payload),].join("\n")
  puts
  puts [full_line_label("winnowed"), winnowed_payload].join("\n")
end
